import optparse, os
from collections import defaultdict
from dbwrap import *

from sage.all import DiGraph

def doctest_graph(db, source):
    latest = list(db.execute("select run_id from timings where source = ? order by run_id limit 1", source))
    if not latest:
        raise ValueError, "No such doctest: '%s'" % source
    run_id = latest[0]['run_id']
    filter_clause = "type='example' and run_id=? and doctest in (select doctest from timings where source=? and run_id=?)"
    filter_args = run_id, source, run_id
    counts = db.execute("select count(*) from timings where source=? and " + filter_clause, source, *filter_args).next()["count(*)"]
    print counts, "occurrences"
    states = len(list(db.execute("select running_state from timings where source=? and " + filter_clause + " group by running_state", source, *filter_args)))
    print states, "initial states"
    query = "select doctest, sequence_number, source, predecessors from timings where " + filter_clause
    examples = dict()
    successors = defaultdict(list)
    predecessors = defaultdict(list)
    all = db.execute(query, *filter_args)
    for example in all:
        if example['source'] == source and False:
            print example
        doctest = example['doctest']
        num = example['sequence_number']
        examples[doctest, num] = example['source']
        for predecessor in example['predecessors']:
            predecessors[doctest, num].append((doctest, predecessor))
            successors[doctest, predecessor].append((doctest, num))
    G = DiGraph()
    def extend_forward(node):
        for s in successors[node]:
            G.add_edge(examples[node], examples[s])
            extend_forward(s)
    def extend_backward(node):
        for s in predecessors[node]:
            G.add_edge(examples[s], examples[node])
            extend_backward(s)
    for example in db.execute("select doctest, sequence_number from timings where run_id=? and source=?", run_id, source):
        node = (example['doctest'], example['sequence_number'])
        extend_forward(node)
        extend_backward(node)
    return G
    
    
def print_graph(G, root, indent="", inverted=None):
    if inverted is None:
        # Print both
        print_graph(G.reverse(), root, indent, inverted=True)
        print
        print_graph(G, root, indent, inverted=False)
        return
    if not inverted:
        print indent, root.strip()
    for node in G.neighbor_out_iterator(root):
        print_graph(G, node, "    " + indent, inverted)
    if inverted:
        print indent, root.strip()


if __name__ == "__main__":

    parser = optparse.OptionParser()
    parser.add_option("-d", "--database", dest="database", default="timings.db", metavar="/path/to/db.sqlite")
    parser.add_option("-r", "--run_id", dest="run_id")
    parser.add_option("-o", "--output", dest="output", default="doctest")
    (options, args) = parser.parse_args()

    assert len(args) == 1
    source = args[0]
    db = DBWrapper(options.database)
    G = doctest_graph(db, source + "\n")
    print_graph(G, source + "\n")
    G.save(options.output + ".sobj")
    G.plot().plot(figsize=50).save(options.output + ".png")
    G.graphviz_to_file_named(options.output + ".dot")
    os.system("unset LIBRARY_PATH; unset LD_LIBRARY_PATH; unset DYLD_LIBRARY_PATH; dot -Tpng -O " + options.output + ".dot")
    
